import argparse
import numpy as np
import torch
import triton
import triton.language as tl

@triton.jit
def l2norm_fwd_kernel(
    x,
    y,
    rstd,
    eps,
    T: tl.constexpr,
    D: tl.constexpr,
    BD: tl.constexpr,
    NB: tl.constexpr,
    BT: tl.constexpr,
):
    i_t = tl.program_id(0)
    p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
    p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
    p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))

    b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
    b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps)
    b_y = b_x * b_rstd[:, None]

    tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))

def test_l2norm_fwd_kernel(output_file):
    # 设置随机种子以确保可重复性
    torch.manual_seed(42)

    # 定义参数
    B = 2        # 批量大小
    T = 64       # 序列长度
    D = 32       # 特征维度
    BD = 8       # block大小 for D
    NB = 4       # block大小 for N
    BT = 16      # block大小 for T

    # 生成随机输入张量
    device = 'npu'  # 使用NPU设备
    dtype = torch.float32  # 使用float32以匹配内核的内部类型

    # 输入张量
    x = torch.randn(B, T, D, dtype=dtype, device=device)
    y = torch.randn(B, T, D, dtype=dtype, device=device)
    rstd = torch.randn(B, T, dtype=dtype, device=device)
    eps = 1e-6  # 防止除零

    # 计算网格大小
    num_blocks_t = triton.cdiv(T, BT)
    grid = (num_blocks_t,)

    # 调用内核函数
    l2norm_fwd_kernel[grid](
        x, y, rstd, eps,
        T=T, D=D, BD=BD, NB=NB, BT=BT
    )

    y_numpy = y.cpu().detach().numpy()
    np.savetxt(output_file, y_numpy.reshape(-1, y_numpy.shape[-1]))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test Causal Conv1D Update Kernel')
    parser.add_argument('--output', type=str, default='default_output.txt', 
                        help='Output file name (default: default_output.txt)')
    args = parser.parse_args()
    test_l2norm_fwd_kernel(args.output)